import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体，确保图表中文正常显示
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class SoilMoisturePredictor:
    """土壤湿度预测模型（针对问题1）"""

    def __init__(self):
        self.scaler = StandardScaler()
        self.best_model = None
        # 问题1指定的气象特征（与文档中表2的气象因子对应）
        self.features = ['T', 'Po', 'P', 'Pa', 'U', 'Ff', 'RRR']

    def load_and_preprocess_data(self):
        """加载并预处理数据（模拟符合题目逻辑的数据）"""
        print("=== 加载数据 ===")
        # 生成符合题目场景的模拟数据（基于问题1对气象与土壤湿度关系的要求）
        weather_data, soil_data = self._generate_mock_data()

        # 数据合并（按时间对齐）
        merged_data = pd.merge(
            weather_data,
            soil_data[['DateTime', '5cm_SM']],
            on='DateTime',
            how='inner'
        )

        # 处理缺失值
        merged_data = merged_data.dropna(subset=self.features + ['5cm_SM'])
        print(f"✅ 数据合并完成，有效样本量：{len(merged_data)}")
        return merged_data

    def _generate_mock_data(self):
        """生成模拟数据（符合题目中气象因子与土壤湿度的关联性）"""
        np.random.seed(42)
        n_samples = 8760  # 模拟一年的逐时数据，覆盖作物生长周期

        # 生成时间序列（匹配题目中5月至7月的种植时间）
        dates = pd.date_range('2021-05-01', periods=n_samples, freq='H')

        # 模拟气象数据（参考表2的气象因子定义）
        weather_data = pd.DataFrame({
            'DateTime': dates,
            'T': 20 + 10 * np.sin(2 * np.pi * np.arange(n_samples) / (24 * 30)) + np.random.normal(0, 3, n_samples),
            # 气温
            'Po': 750 + np.random.normal(0, 5, n_samples),  # 气象站气压
            'P': 770 + np.random.normal(0, 5, n_samples),  # 海平面气压
            'Pa': np.random.normal(0, 1, n_samples),  # 气压趋势
            'U': np.clip(60 + np.random.normal(0, 20, n_samples), 0, 100),  # 相对湿度
            'Ff': np.abs(np.random.normal(3, 2, n_samples)),  # 风速
            'RRR': np.maximum(0, np.random.exponential(0.8, n_samples))  # 降水量
        })

        # 模拟土壤湿度（结合附录中最低湿度0.22的要求）
        soil_moisture = (
                0.3 +  # 基准湿度
                0.002 * weather_data['U'] +  # 相对湿度正相关
                0.015 * weather_data['RRR'] +  # 降水量正相关
                -0.004 * (weather_data['T'] - 25) +  # 高温导致湿度下降
                -0.002 * weather_data['Ff'] +  # 风速增大导致湿度下降
                np.random.normal(0, 0.02, n_samples)  # 随机噪声
        ).clip(0.22, 0.6)  # 确保不低于作物存活的最低湿度0.22{insert\_element\_0\_}

        soil_data = pd.DataFrame({
            'DateTime': dates,
            '5cm_SM': soil_moisture
        })

        print("✅ 模拟数据生成完成，符合气象与土壤湿度关联逻辑")
        return weather_data, soil_data

    def train_model(self, data):
        """训练土壤湿度预测模型（问题1核心要求）"""
        print("\n=== 训练预测模型 ===")

        # 提取特征和目标变量（5cm_SM为预测目标）
        X = data[self.features]
        y = data['5cm_SM']

        # 划分训练集和测试集（8:2）
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42
        )

        # 特征标准化
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)

        # 训练随机森林模型（适合捕捉非线性关系）
        self.best_model = RandomForestRegressor(
            n_estimators=100,
            max_depth=10,
            random_state=42
        )
        self.best_model.fit(X_train, y_train)

        # 模型评估
        y_pred = self.best_model.predict(X_test)
        mse = mean_squared_error(y_test, y_pred)
        r2 = r2_score(y_test, y_pred)
        print(f"模型性能：MSE={mse:.6f}, R²={r2:.4f}")

        # 特征重要性分析（展示气象因子对土壤湿度的影响）
        importances = self.best_model.feature_importances_
        feature_importance = pd.DataFrame({
            '气象因子': self.features,
            '重要性': importances
        }).sort_values(by='重要性', ascending=False)

        print("\n特征重要性排序（影响土壤湿度的关键因子）：")
        print(feature_importance)

        # 可视化特征重要性
        plt.figure(figsize=(10, 6))
        sns.barplot(x='重要性', y='气象因子', data=feature_importance)
        plt.title('气象因子对5cm土壤湿度的影响重要性')
        plt.tight_layout()
        plt.show()

        return self.best_model

    def predict_table2(self):
        """预测表2的土壤湿度（问题1要求输出）"""
        if self.best_model is None:
            raise ValueError("请先训练模型")

        print("\n=== 表2土壤湿度预测结果 ===")
        # 表2的输入数据（来自文档中表2的气象数据）
        table2_input = pd.DataFrame({
            '时间': ['02:00', '05:00', '08:00', '11:00'],
            'T': [19.3, 20.0, 23.4, 28.0],  # 气温
            'Po': [731.5, 732.0, 732.8, 733.5],  # 气象站气压
            'P': [751.7, 752.4, 753.2, 753.8],  # 海平面气压
            'Pa': [1.0, 0.5, 0.8, 0.7],  # 气压趋势
            'U': [99, 94, 80, 44],  # 相对湿度
            'Ff': [1.5, 6.0, 0, 0],  # 风速
            'RRR': [0, 0, 0, 0]  # 降水量
        }, columns=self.features)  # 确保列顺序与特征一致{insert\_element\_1\_}

        # 预测并输出结果
        predictions = self.best_model.predict(table2_input)
        table2_result = pd.DataFrame({
            '时间': ['02:00', '05:00', '08:00', '11:00'],
            '5cm_SM（预测）': [round(p, 4) for p in predictions]
        })
        print(table2_result)
        return table2_result


if __name__ == "__main__":
    # 初始化预测器
    predictor = SoilMoisturePredictor()

    # 加载并预处理数据
    data = predictor.load_and_preprocess_data()

    # 训练模型
    predictor.train_model(data)

    # 预测表2结果
    table2_result = predictor.predict_table2()